# -*- encoding: utf-8 -*-
"""
@author: aoboxia
@date: 2022/10/31
"""
from datasets2 import *
import torch
from tqdm import tqdm
device = torch.device('cuda:0')


def evaluate(model, test_dataloader, record_path=None, use_mask_logit=False, id2label=None):
    if record_path:
        f = open(record_path, 'w', encoding='utf8')
    else:
        f = None

    def _precision_recall_f1(right, predict, total):
        """
        :param right: int, the count of right prediction
        :param predict: int, the count of prediction
        :param total: int, the count of labels
        :return: p(precision, Float), r(recall, Float), f(f1_score, Float)
        """
        p, r, f = 0.0, 0.0, 0.0
        if predict > 0:
            p = float(right) / predict
        if total > 0:
            r = float(right) / total
        if p + r > 0:
            f = p * r * 2 / (p + r)
        return p, r, f

    def calc_char(crf_output, slot_values, max_len):
        crf_output = crf_output[0]
        slot_values = torch.squeeze(slot_values)
        right_total, predict_total, gold_total = 0, 0, 0

        for i in range(1, max_len):
            type1 = int((slot_values[i] + 1) / 2)
            type2 = int((crf_output[i] + 1) / 2)
            if slot_values[i] != 0:
                gold_total += 1
                if crf_output[i] == slot_values[i]:
                    right_total += 1

            if crf_output[i] != 0:
                predict_total += 1

        return right_total, predict_total, gold_total

    def get_all_slots(crf_output, tokens):
        res = []
        tmp = ""
        type = -1
        for i in range(len(crf_output)):
            if crf_output[i] == 0:
                if tmp != "" and type != -1:
                    res.append((type, tmp))
                    type = -1
                    tmp = ""
            else:
                tmp += tokens[i]
                type = int((crf_output[i] + 1) / 2)

        if tmp != "" and type != -1:
            res.append((type, tmp))
            type = -1
            tmp = ""

        return res


    right_count_list = [0 for _ in range(len(id2label))]
    gold_count_list = [0 for _ in range(len(id2label))]
    predicted_count_list = [0 for _ in range(len(id2label))]
    precision_dict = dict()
    recall_dict = dict()
    fscore_dict = dict()
    right_total, predict_total, gold_total = 0, 0, 0

    char_right_count_list = [0 for _ in range(len(id2label))]
    char_gold_count_list = [0 for _ in range(len(id2label))]
    char_predicted_count_list = [0 for _ in range(len(id2label))]
    char_precision_dict = dict()
    char_recall_dict = dict()
    char_fscore_dict = dict()
    char_right_total, char_predict_total, char_gold_total = 0, 0, 0
    key = 1.3

    model.eval()
    test_dataloader_tqdm = tqdm(test_dataloader)
    for tokenized_ids, segment_ids, slot_values, mask_logit in test_dataloader_tqdm:
        tokenized_ids = tokenized_ids.long().to(device)
        segment_ids = segment_ids.long().to(device)
        mask_logit = mask_logit.long().to(device)
        inputs = (tokenized_ids, segment_ids)

        if use_mask_logit:
            ner_output, crf_output = model(inputs, mask_logit=mask_logit)
        else:
            ner_output, crf_output = model(inputs)

        mask = (tokenized_ids != 0)
        max_len = sum(sum(mask))
        # print(mask)
        # print(max_len)

        # print(tokenized_ids[0])
        tokens = ids2tokens(tokenized_ids[0])
        # print(slot_values)
        # print("----------")
        # print(crf_output)
        target_values = get_all_slots(slot_values[0], tokens)
        predcit_values = get_all_slots(crf_output[0], tokens)
        # print(target_values, predcit_values)

        tokens = [t for t in tokens if t != '[PAD]']
        f.write(''.join(tokens[1:]))
        f.write('\t\t\t')
        f.write(str(target_values))
        f.write('\t\t\t')
        f.write(str(predcit_values))
        f.write('\n')

        # count for the gold and right items
        for gold in target_values:
            gold_count_list[gold[0]] += 1
            for label in predcit_values:
                if gold[0] == label[0] and gold[1] == label[1]:
                    right_count_list[gold[0]] += 1

        # count for the predicted items
        for label in predcit_values:
            predicted_count_list[label[0]] += 1


        char_right, char_predict, char_gold = calc_char(crf_output, slot_values, max_len)

        char_right_total += char_right
        char_predict_total += char_predict
        char_gold_total += char_gold
        precision, recall, f1 = _precision_recall_f1(char_right_total, char_predict_total, char_gold_total)

        # test_dataloader_tqdm.set_description(
        #     "char_precision: {:3f} , char_recall: {:3f} , char_f1: {:3f} ".format(precision, recall, f1))
        # test_dataloader_tqdm.set_description(
        #     "word_precision: {:3f} , word_recall: {:3f} , word_f1: {:3f} char_precision: {:3f} , char_recall: {:3f} , char_f1: {:3f} ".format(
        #         precision_micro, recall_micro, micro_f1, precision, recall, f1))

    for i, label in id2label.items():
        label = label + '_' + str(i)
        precision_dict[label], recall_dict[label], fscore_dict[label] = _precision_recall_f1(right_count_list[i],
                                                                                             predicted_count_list[
                                                                                                 i],
                                                                                             gold_count_list[i])

        print(label, precision_dict[label], recall_dict[label], fscore_dict[label])

        right_total += right_count_list[i]
        gold_total += gold_count_list[i]
        predict_total += predicted_count_list[i]

    # Macro-F1
    precision_macro = sum([v for _, v in precision_dict.items()]) / len(list(precision_dict.keys()))
    recall_macro = sum([v for _, v in recall_dict.items()]) / len(list(precision_dict.keys()))
    macro_f1 = sum([v for _, v in fscore_dict.items()]) / len(list(fscore_dict.keys()))
    # Micro-F1
    precision_micro = float(right_total) / predict_total if predict_total > 0 else 0.0
    recall_micro = float(right_total) / gold_total
    micro_f1 = 2 * precision_micro * recall_micro / (precision_micro + recall_micro) if (precision_micro + recall_micro) > 0 else 0.0

    precision, recall, f1 = _precision_recall_f1(char_right_total, char_predict_total, char_gold_total)

    # print(precision_macro, recall_macro, macro_f1)
    # print(precision_micro, recall_micro, micro_f1)

    return precision_macro, recall_macro, macro_f1, precision, recall, f1